import streamlit as st
import wave
import json
from vosk import Model, KaldiRecognizer
from pydub import AudioSegment
from pydub.playback import play
import io
import os
import sounddevice as sd
import soundfile as sf
import numpy as np

# 加载Vosk模型
model = Model("zhongwen")

# Streamlit页面设置
st.title("语音处理客户端")
st.write("这是一个简单的语音处理客户端，支持录音、下载音频、转成中文等功能。")

# 录音功能
def record_audio(filename, duration):
    fs = 16000  # 采样率
    st.write(f"开始录音，时长：{duration}秒...")
    recording = sd.rec(int(duration * fs), samplerate=fs, channels=1, dtype='int16')
    sd.wait()  # 等待录音完成
    sf.write(filename, recording, fs)
    st.write("录音结束。")

# 下载音频文件
def save_audio(file):
    with open(file, "rb") as f:
        st.download_button(
            label="下载音频",
            data=f,
            file_name=file,
            mime="audio/wav"
        )

# 语音识别功能
def transcribe_audio(file):
    try:
        wf = wave.open(file, "rb")
        if wf.getnchannels() != 1 or wf.getsampwidth() != 2 or wf.getframerate() != 16000:
            st.error("音频文件必须是单声道、16位、16kHz的WAV格式。")
            return

        rec = KaldiRecognizer(model, 16000)
        results = []
        while True:
            data = wf.readframes(4000)
            if len(data) == 0:
                break
            if rec.AcceptWaveform(data):
                result = json.loads(rec.Result())
                results.append(result['text'])

        final_result = json.loads(rec.FinalResult())
        results.append(final_result['text'])
        transcription = " ".join(results)
        return transcription
    except wave.Error as e:
        st.error(f"无法读取音频文件：{e}")
        return None

# 主程序
def main():
    # 录音功能
    st.write("### 录音设置")
    duration = st.number_input("请输入录音时长（秒）", min_value=1, max_value=60, value=5)  # 用户输入录音时长

    if st.button("开始录音"):
        audio_file = "recorded_audio.wav"
        record_audio(audio_file, duration)
        st.audio(audio_file, format='audio/wav')

    # 上传音频文件
    uploaded_file = st.file_uploader("上传音频文件", type=["wav"])
    if uploaded_file is not None:
        audio_file = "uploaded_audio.wav"
        with open(audio_file, "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.audio(audio_file, format='audio/wav')

    # 下载音频
    if os.path.exists("recorded_audio.wav"):
        save_audio("recorded_audio.wav")
    if os.path.exists("uploaded_audio.wav"):
        save_audio("uploaded_audio.wav")

    # 语音识别
    if st.button("转成中文"):
        if os.path.exists("recorded_audio.wav"):
            transcription = transcribe_audio("recorded_audio.wav")
            if transcription:
                st.write("识别结果：", transcription)
        elif os.path.exists("uploaded_audio.wav"):
            transcription = transcribe_audio("uploaded_audio.wav")
            if transcription:
                st.write("识别结果：", transcription)
        else:
            st.error("请先录音或上传音频文件。")

if __name__ == "__main__":
    main()